// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "ppl/nn/oputils/onnx/reshape_split.h"
#include "ppl/nn/runtime/tensor_impl.h"
#include "ppl/nn/common/logger.h"
using namespace ppl::common;

namespace ppl { namespace nn { namespace onnx {

RetCode ReshapeSplit(InputOutputInfo* info, const ir::Attr* arg, const int64_t* split_point) {
    auto param = static_cast<const SplitParam*>(arg);
    auto split_size = info->GetInput<TensorImpl>(1)->GetShape()->GetDim(0);
    const TensorShape& shape = *info->GetInput<TensorImpl>(0)->GetShape();
    int dim_count = shape.GetDimCount();
    uint32_t split_axis = param->axis >= 0 ? param->axis : param->axis + dim_count;
    if (split_size == 0) {
        uint32_t in_dim = shape.GetDim(split_axis);
        uint32_t split_dim = in_dim / info->GetOutputCount();
        if (split_dim * info->GetOutputCount() != in_dim) {
            LOG(DEBUG) << "ERROR: split_dim[" << split_dim << "] * output_count[" << info->GetOutputCount()
                       << "] != in_dim[" << in_dim << "]";
            return RC_INVALID_VALUE;
        }

        std::vector<int64_t> tmp_output_dim(dim_count);
        for (uint32_t i = 0; i < info->GetOutputCount(); ++i) {
            for (int it = 0; it < dim_count; ++it) {
                tmp_output_dim[it] = shape.GetDim(it);
            }
            tmp_output_dim[split_axis] = split_dim;
            info->GetOutput<TensorImpl>(i)->GetShape()->Reshape(tmp_output_dim);
        }
        return RC_SUCCESS;
    }
    if (info->GetOutputCount() != split_size) {
        LOG(DEBUG) << "ERROR: output_count[" << info->GetOutputCount() << "] != split point size["
                   << split_size << "]";
        return RC_INVALID_VALUE;
    }

    uint32_t sum_dim = 0;
    for (uint32_t i = 0; i < split_size; ++i) {
        sum_dim += split_point[i];
    }
    if (sum_dim != shape.GetDim(split_axis)) {
        LOG(DEBUG) << "ERROR: sum_dim[" << sum_dim << "] != dim[" << shape.GetDim(split_axis) << "] of axis["
                   << split_axis << "]";
        return RC_INVALID_VALUE;
    }

    std::vector<int64_t> tmp_output_dim(dim_count);
    for (uint32_t i = 0; i < split_size; ++i) {
        for (int it = 0; it < dim_count; ++it) {
            tmp_output_dim[it] = shape.GetDim(it);
        }
        tmp_output_dim[split_axis] = split_point[i];
        info->GetOutput<TensorImpl>(i)->GetShape()->Reshape(tmp_output_dim);
    }

    return RC_SUCCESS;
}

RetCode ReshapeSplit(InputOutputInfo* info, const ir::Attr* arg) {
    if (info->GetInputCount() > 2) {
        LOG(DEBUG) << "ERROR: input count[" << info->GetInputCount() << "] != 2.";
        return RC_INVALID_VALUE;
    }

    if (info->GetInputCount() == 1) {
        auto in_shape = *info->GetInput<TensorImpl>(0)->GetShape();
        info->GetOutput<TensorImpl>(0)->GetShape()->Reshape(in_shape.GetDims(), in_shape.GetRealDimCount());
        return RC_SUCCESS;
    }

    auto shape_data = info->GetInput<TensorImpl>(1)->GetBufferPtr<int64_t>();
    return ReshapeSplit(info, arg, shape_data);
}

}}} // namespace ppl::nn::onnx
